use core::{
    alloc::{AllocError, Allocator, Layout},
    borrow::BorrowMut,
    cell::RefCell,
    iter::Map,
    marker::PhantomData,
    mem::size_of,
    ops::{Index, IndexMut},
    ptr::{slice_from_raw_parts_mut, NonNull},
    slice::{ChunksExact, ChunksExactMut},
};

use alloc::{alloc::Global, boxed::Box};

struct FixedRaw<MetaType, A: Allocator> {
    raw: Box<[usize], A>,
    phantom: PhantomData<MetaType>,
}

impl<MetaType, A: Allocator> IndexMut<usize> for FixedRaw<MetaType, A> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        self.get_mut(index).unwrap()
    }
}
impl<MetaType, A: Allocator> Index<usize> for FixedRaw<MetaType, A> {
    type Output = usize;

    fn index(&self, index: usize) -> &Self::Output {
        self.get(index).unwrap()
    }
}

impl<'a, MetaType, A: Allocator> IntoIterator for &'a mut FixedRaw<MetaType, A> {
    type Item = &'a mut usize;

    type IntoIter = Map<ChunksExactMut<'a, usize>, fn(&'a mut [usize]) -> &'a mut usize>;

    fn into_iter(self) -> Self::IntoIter {
        self.raw
            .chunks_exact_mut(FixedRaw::<MetaType, A>::container_num())
            .into_iter()
            .map(|x| &mut x[0])
    }
}
impl<'a, MetaType, A: Allocator> IntoIterator for &'a FixedRaw<MetaType, A> {
    type Item = &'a usize;

    type IntoIter = Map<ChunksExact<'a, usize>, fn(&'a [usize]) -> &'a usize>;

    fn into_iter(self) -> Self::IntoIter {
        self.raw
            .chunks_exact(FixedRaw::<MetaType, A>::container_num())
            .into_iter()
            .map(|x| &x[0])
    }
}

impl<MetaType, A: Allocator> FixedRaw<MetaType, A> {
    fn get(&self, index: usize) -> Option<&usize> {
        Some(
            &self
                .raw
                .chunks_exact(Self::container_num())
                .into_iter()
                .nth(index)?[0],
        )
    }
    fn get_mut(&mut self, index: usize) -> Option<&mut usize> {
        Some(
            &mut self
                .raw
                .chunks_exact_mut(Self::container_num())
                .into_iter()
                .nth(index)?[0],
        )
    }
    fn container_num() -> usize {
        if size_of::<MetaType>() % size_of::<usize>() != 0 {
            return (size_of::<MetaType>() + size_of::<usize>()) / size_of::<usize>();
        }
        size_of::<MetaType>() / size_of::<usize>()
    }
    fn new_in(len: usize, alloc: A) -> Result<Self, AllocError> {
        if size_of::<MetaType>() < size_of::<usize>() {
            return Err(AllocError);
        }
        let mut ret = Self {
            raw: unsafe {
                Box::<[usize], A>::new_uninit_slice_in(len * Self::container_num(), alloc)
                    .assume_init()
            },
            phantom: PhantomData,
        };
        ret.borrow_mut().into_iter().fold(1, |acc, ele| {
            *ele = acc;
            acc + 1
        });
        if let Some(last) = ret.borrow_mut().into_iter().last() {
            *last = 0;
        }
        Ok(ret)
    }
    fn pop<'a>(&'a mut self) -> Option<&'a mut MetaType> {
        let first_head_pos = self[0];
        if first_head_pos == 0 {
            return None;
        }
        let second_head_pos = *self.get(first_head_pos)?;
        self[0] = second_head_pos;
        let first_head = self.get_mut(first_head_pos)?;
        Some(unsafe { (first_head as *mut usize).cast::<MetaType>().as_mut() }?)
    }
    fn push(&mut self, target: &mut MetaType) {
        let target_head = (target as *mut MetaType).cast::<usize>();
        let target_head_pos =
            unsafe { target_head.offset_from(self.raw.as_ptr()) } as usize / Self::container_num();
        if self.into_iter().len() > target_head_pos {
            let first_head_pos = self[0];
            unsafe {
                *target_head = first_head_pos;
            }
            self[0] = target_head_pos;
        }
    }
}

pub struct Fixed<Meta, A: Allocator = Global> {
    raw: RefCell<FixedRaw<Meta, A>>,
}

unsafe impl<MetaType, A: Allocator> Allocator for Fixed<MetaType, A> {
    fn allocate(
        &self,
        layout: core::alloc::Layout,
    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
        if layout != Layout::new::<MetaType>() {
            return Err(AllocError);
        }
        let data = RefCell::borrow_mut(&self.raw)
            .pop()
            .map_or_else(|| Err(AllocError), |x| Ok(x as *mut MetaType))?
            .cast::<u8>();
        let slice_ptr = slice_from_raw_parts_mut(data, size_of::<MetaType>() / size_of::<u8>());
        Ok(unsafe { NonNull::new_unchecked(slice_ptr) })
    }

    unsafe fn deallocate(&self, ptr: core::ptr::NonNull<u8>, layout: core::alloc::Layout) {
        if layout == Layout::new::<MetaType>() {
            let target = ptr.cast::<MetaType>().as_mut();
            RefCell::borrow_mut(&self.raw).push(target);
        }
    }
}

impl<MetaType, A: Allocator> Fixed<MetaType, A> {
    /// 注意如果 MetaType 的大小不等于 usize，则会 n 个 usize 大小的空间，使其能够正好覆盖 MetaType
    pub fn new_in(len: usize, alloc: A) -> Result<Self, AllocError> {
        Ok(Self {
            raw: RefCell::new(FixedRaw::new_in(len, alloc)?),
        })
    }
}

impl<MetaType> Fixed<MetaType> {
    /// 注意如果 MetaType 的大小不等于 usize，则会 n 个 usize 大小的空间，使其能够正好覆盖 MetaType
    pub fn new(len: usize) -> Result<Self, AllocError> {
        Ok(Self::new_in(len, Global)?)
    }
}
